"""
导入大模型岗位数据命令
"""
import pandas as pd
from django.core.management.base import BaseCommand
from django.conf import settings
from apps.positions.models import AIPosition
from utils.data_clean import clean_salary, split_location, normalize_education, clean_text_field, validate_salary


class Command(BaseCommand):
    help = '导入大模型岗位CSV数据到数据库'

    def handle(self, *args, **options):
        self.stdout.write('='* 50)
        self.stdout.write('开始导入大模型岗位数据...')
        self.stdout.write('='* 50)

        # CSV文件路径
        csv_path = settings.DATA_DIR / settings.AI_JOBS_CSV_FILE
        
        if not csv_path.exists():
            self.stdout.write(self.style.ERROR(f'CSV文件不存在: {csv_path}'))
            return

        try:
            # 读取CSV
            self.stdout.write(f'正在读取CSV文件: {csv_path}')
            df = pd.read_csv(csv_path, encoding='utf-8')
            self.stdout.write(self.style.SUCCESS(f'成功读取 {len(df)} 条数据'))
        except Exception as e:
            self.stdout.write(self.style.ERROR(f'读取CSV失败: {e}'))
            return

        # 数据清洗
        self.stdout.write('开始数据清洗...')
        df = self.clean_data(df)
        self.stdout.write(self.style.SUCCESS('数据清洗完成'))

        # 导入数据库
        self.stdout.write('开始导入数据库...')
        positions = []
        success_count = 0
        error_count = 0
        
        for idx, row in df.iterrows():
            try:
                position = AIPosition(
                    position_name=row['岗位名称'],
                    location=row.get('工作地点'),
                    location_city=row.get('location_city'),
                    location_district=row.get('location_district'),
                    salary=row.get('岗位薪资'),
                    salary_min=row.get('salary_min'),
                    salary_max=row.get('salary_max'),
                    salary_avg=row.get('salary_avg'),
                    salary_months=row.get('salary_months'),
                    experience=row.get('经验要求'),
                    education=row.get('学历要求'),
                    education_normalized=row.get('education_normalized'),
                    tags=row.get('岗位标签'),
                    company_name=row.get('企业名称'),
                    company_industry=row.get('企业行业'),
                    company_scale=row.get('企业规模'),
                    financing_status=row.get('融资状况'),
                )
                positions.append(position)
                success_count += 1

                # 批量插入（每500条）
                if len(positions) >= 500:
                    AIPosition.objects.bulk_create(positions, ignore_conflicts=True)
                    self.stdout.write(f'已导入 {success_count} 条数据')
                    positions = []

            except Exception as e:
                error_count += 1
                if error_count <= 5:  # 只显示前5个错误
                    self.stdout.write(self.style.WARNING(f'第 {idx + 1} 行导入失败: {e}'))

        # 插入剩余数据
        if positions:
            AIPosition.objects.bulk_create(positions, ignore_conflicts=True)

        self.stdout.write('='* 50)
        self.stdout.write(self.style.SUCCESS(
            f'数据导入完成！成功: {success_count}, 失败: {error_count}'
        ))
        self.stdout.write('='* 50)

    def clean_data(self, df):
        """数据清洗"""
        # 1. 解析薪资（包含月数）
        df['salary_min'] = None
        df['salary_max'] = None
        df['salary_avg'] = None
        df['salary_months'] = None
        
        for idx, row in df.iterrows():
            salary_str = row.get('岗位薪资')
            if pd.notna(salary_str):
                salary_min, salary_max, salary_avg, salary_months = clean_salary(salary_str)
                if validate_salary(salary_min, salary_max):
                    df.at[idx, 'salary_min'] = salary_min
                    df.at[idx, 'salary_max'] = salary_max
                    df.at[idx, 'salary_avg'] = salary_avg
                    df.at[idx, 'salary_months'] = salary_months
        
        # 2. 拆分地点
        df['location_city'] = None
        df['location_district'] = None
        
        for idx, row in df.iterrows():
            location_str = row.get('工作地点')
            if pd.notna(location_str):
                city, district = split_location(location_str)
                df.at[idx, 'location_city'] = city
                df.at[idx, 'location_district'] = district
        
        # 3. 规范化学历
        if '学历要求' in df.columns:
            df['education_normalized'] = df['学历要求'].apply(normalize_education)
        
        # 4. 处理"Unknown"标签
        if '岗位标签' in df.columns:
            df.loc[df['岗位标签'] == 'Unknown', '岗位标签'] = None
        
        # 5. 清理文本字段
        text_fields = ['岗位名称', '企业名称', '企业行业', '企业规模']
        for field in text_fields:
            if field in df.columns:
                df[field] = df[field].apply(clean_text_field)
        
        # 6. 去除岗位名称为空的记录
        df = df.dropna(subset=['岗位名称'])
        
        return df

